#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HF 模型下载 GUI 工具（自动使用 hf-mirror.com 镜像）
支持 repo_id 输入、类型选择、自动路径、日志输出（含 snapshot_download 实时日志）

repo_id 示例：speechbrain/spkrec-ecapa-voxceleb

官方推荐方法：
# Windows CMD
set HF_ENDPOINT=https://hf-mirror.com  
hf download ^
  --repo-type model ^
  speechbrain/spkrec-ecapa-voxceleb ^
  --local-dir models/speechbrain/spkrec-ecapa-voxceleb ^
"""
import os
import sys
import threading
import tkinter as tk
from tkinter import ttk, scrolledtext
from pathlib import Path
import importlib.util
import logging
from queue import Queue, Empty
from huggingface_hub import snapshot_download
import time

# ==============================
# 0. 修复 sys.stdout / sys.stderr 为 None 的问题
# ==============================
class DummyWriter:
    def write(self, data):
        pass
    def flush(self):
        pass

if sys.stdout is None:
    sys.stdout = DummyWriter()
if sys.stderr is None:
    sys.stderr = DummyWriter()

# ==============================
# 1. 环境配置
# ==============================
HF_MIRROR_URL = "https://hf-mirror.com"
if os.environ.get("HF_ENDPOINT", "").strip() != HF_MIRROR_URL:
    os.environ["HF_ENDPOINT"] = HF_MIRROR_URL
# 自动启用 hf-transfer（如果已安装）
if importlib.util.find_spec("hf_transfer") is not None:
    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
# 获取系统 HF_TOKEN（如果存在）
HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None

# ==============================
# 2. 日志转发 Handler（用于捕获 huggingface_hub 日志）
# ==============================
class GuiLogHandler(logging.Handler):
    def __init__(self, log_queue):
        super().__init__()
        self.log_queue = log_queue

    def emit(self, record):
        msg = self.format(record)
        self.log_queue.put(msg)

# ==============================
# 3. 下载函数（在线程中运行）
# ==============================
def download_model(repo_id, repo_type, token, log_widget, on_complete_callback):
    log_queue = Queue()

    def safe_log(msg):
        log_widget.after(0, lambda: log_widget.insert(tk.END, msg + "\n"))
        log_widget.after(0, lambda: log_widget.see(tk.END))

    # 日志消费线程：持续从队列读取并更新 GUI
    def log_consumer():
        while not stop_log_consumer.is_set():
            try:
                msg = log_queue.get(timeout=0.1)
                safe_log(msg)
            except Empty:
                continue

    if not repo_id.strip():
        safe_log("❌ 错误：repo_id 不能为空！")
        on_complete_callback()
        return

    safe_path = repo_id.replace("/", os.sep)
    script_dir = Path(__file__).parent.resolve()
    local_dir = script_dir / "models" / safe_path
    local_dir.mkdir(parents=True, exist_ok=True)
    safe_log(f"📥 开始下载: {repo_id}")
    safe_log(f"📁 本地路径: {local_dir}")

    # 设置日志捕获
    logger = logging.getLogger("huggingface_hub")
    handler = GuiLogHandler(log_queue)
    formatter = logging.Formatter('[HF] %(levelname)s: %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

    # 启动日志消费线程
    stop_log_consumer = threading.Event()
    log_thread = threading.Thread(target=log_consumer, daemon=True)
    log_thread.start()

    try:
        snapshot_download(
            repo_id=repo_id,
            repo_type=repo_type,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,
            token=token,
        )
        safe_log("✅ 下载成功！")
    except Exception as e:
        safe_log(f"❌ 下载失败: {e}")
    finally:
        # 停止日志消费线程
        stop_log_consumer.set()
        log_thread.join(timeout=1)  # 等待最多1秒
        # 清空可能残留的日志
        while not log_queue.empty():
            try:
                msg = log_queue.get_nowait()
                safe_log(msg)
            except Empty:
                break
        logger.removeHandler(handler)
        on_complete_callback()

# ==============================
# 4. GUI 主窗口
# ==============================
class HFDownloadGUI:
    def __init__(self, root):
        self.root = root
        root.title("HF 模型下载工具 (hf-mirror.com)")
        root.geometry("650x550")
        root.resizable(True, True)

        # Repo ID 输入
        ttk.Label(root, text="Repo ID:").grid(row=0, column=0, padx=10, pady=10, sticky="w")
        self.repo_id_var = tk.StringVar()
        ttk.Entry(root, textvariable=self.repo_id_var, width=60).grid(row=0, column=1, padx=5, pady=10, sticky="ew")

        # 类型下拉框
        ttk.Label(root, text="类型:").grid(row=1, column=0, padx=10, pady=5, sticky="w")
        self.repo_type_var = tk.StringVar(value="model")
        repo_types = ["model", "dataset", "space"]
        ttk.Combobox(root, textvariable=self.repo_type_var, values=repo_types, state="readonly", width=20).grid(
            row=1, column=1, padx=5, pady=5, sticky="w"
        )

        # Token 输入（可选）
        ttk.Label(root, text="Token (可选):").grid(row=2, column=0, padx=10, pady=5, sticky="w")
        self.token_var = tk.StringVar()
        ttk.Entry(root, textvariable=self.token_var, width=60).grid(row=2, column=1, padx=5, pady=5, sticky="ew")

        # 下载按钮
        self.download_btn = ttk.Button(root, text="开始下载", command=self.start_download)
        self.download_btn.grid(row=3, column=0, columnspan=2, pady=15)

        # 日志输出框
        ttk.Label(root, text="日志:").grid(row=4, column=0, padx=10, sticky="nw")
        self.log_text = scrolledtext.ScrolledText(root, height=18, state="normal", wrap=tk.WORD)
        self.log_text.grid(row=5, column=0, columnspan=2, padx=10, pady=5, sticky="nsew")
        root.grid_rowconfigure(5, weight=1)
        root.grid_columnconfigure(1, weight=1)

    def start_download(self):
        repo_id = self.repo_id_var.get().strip()
        if not repo_id:
            self.log_text.insert(tk.END, "⚠️ 请输入有效的 Repo ID！\n")
            return
        token = (self.token_var.get().strip() or HF_TOKEN) or None
        repo_type = self.repo_type_var.get()
        self.download_btn.config(state="disabled")
        self.log_text.delete(1.0, tk.END)

        def on_complete():
            self.root.after(0, lambda: self.download_btn.config(state="normal"))

        thread = threading.Thread(
            target=download_model,
            args=(repo_id, repo_type, token, self.log_text, on_complete),
            daemon=True
        )
        thread.start()

# ==============================
# 5. 入口
# ==============================
if __name__ == "__main__":
    root = tk.Tk()
    app = HFDownloadGUI(root)
    root.mainloop()